import dgl
import torch
import numpy as np

def get_dist_mat(points):
    """给定点的位置返回距离矩阵"""
    p = points.copy()
    dist_mat = np.zeros((len(p), len(p)))
    for i in range(len(p)):
        for j in range(len(p)):
            dist_mat[i][j] = np.sqrt(np.sum(np.square(p[i] - p[j])))
    return dist_mat


def get_neighbors(dist_mat, self_loop, num_near):
    """根据距离获取最近num_near个邻居的下标"""
    dist_sort = np.argsort(dist_mat, axis=1)
    if self_loop:
        neighbors = dist_sort[:, :num_near]
    else:
        neighbors = dist_sort[:, 1:num_near + 1]
    return neighbors


def tour_to_set(tour):
    """路由转换为边的集合"""
    n = len(tour)
    edge_set = []
    for i in range(n - 1):
        edge_set.append({tour[i], tour[i + 1]})
    edge_set.append({tour[n - 1], tour[0]})
    return edge_set


def data_process(points, tour, self_loop=False, num_near=None):
    """将点位置和路由转为图和标签"""
    n = len(points)
    # 求距离矩阵
    dist_mat = get_dist_mat(points)
    # 将路由转为边的集合
    edge_set = tour_to_set(tour)
    # 根据距离获取最近num_near个邻居的下标
    if num_near is None:
        # 默认是全部节点
        num_near = len(dist_mat)
    neighbors = get_neighbors(dist_mat, self_loop, num_near)
    edges = []  # 边的连接关系
    e_feat = []  # 边的特征
    edge_label = []  # 边的标签
    for i in range(neighbors.shape[0]):
        for j in range(neighbors.shape[1]):
            edges.append([i, neighbors[i][j]])
            e_feat.append(dist_mat[i, neighbors[i][j]])
            if {i, neighbors[i][j]} in edge_set:
                edge_label.append(1)
            else:
                edge_label.append(0)

    edges = torch.tensor(edges, dtype=torch.int64)
    points = torch.tensor(points, dtype=torch.float32)
    e_feat = torch.tensor(e_feat, dtype=torch.float32).view(-1, 1)
    # 数据转为图数据
    graph = dgl.DGLGraph()
    graph.add_nodes(n)
    graph.add_edges(edges[:, 0], edges[:, 1])
    graph.ndata['feat'] = points
    graph.edata['feat'] = e_feat

    edge_label = torch.LongTensor(edge_label)

    return graph, edge_label